package com.ajocer.springbootinit.langchain;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.output.Response;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.springframework.cache.annotation.Cacheable;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.io.IOException;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.text.SimpleDateFormat;
import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/**
 * 讯飞星火嵌入模型适配器
 * 使用讯飞星火专用的Embedding API
 * 注意：这是讯飞文本向量化API，不是大模型API
 */
@Slf4j
public class XingfengEmbeddingModel implements EmbeddingModel {

    private final String appId;
    private final String apiKey;
    private final String apiSecret;
    private final String apiHost;
    private final OkHttpClient httpClient;
    private final String domain;
    
    private static final String DEFAULT_API_HOST = "https://emb-cn-huabei-1.xf-yun.com/";

    /**
     * 使用API参数直接创建XingfengEmbeddingModel
     * 
     * @param apiHost API主机地址，默认为"https://emb-cn-huabei-1.xf-yun.com/"
     * @param appId 应用ID
     * @param apiKey API密钥
     * @param apiSecret API密钥
     */
    public XingfengEmbeddingModel(String apiHost, String appId, String apiKey, String apiSecret) {
        this.appId = appId;
        this.apiKey = apiKey;
        this.apiSecret = apiSecret;
        this.apiHost = apiHost != null && !apiHost.isEmpty() ? apiHost : DEFAULT_API_HOST;
        this.domain = "query"; // 默认为query向量化，用于用户问题
        
        // 配置OkHttpClient
        this.httpClient = new OkHttpClient.Builder()
                .connectTimeout(10, TimeUnit.SECONDS)
                .readTimeout(30, TimeUnit.SECONDS)
                .writeTimeout(30, TimeUnit.SECONDS)
                .build();
        
        log.info("初始化讯飞文本向量化API客户端: apiHost={}, domain={}", this.apiHost, this.domain);
    }
    
    /**
     * 使用API参数直接创建XingfengEmbeddingModel，并指定domain
     * 
     * @param apiHost API主机地址，默认为"https://emb-cn-huabei-1.xf-yun.com/"
     * @param appId 应用ID
     * @param apiKey API密钥
     * @param apiSecret API密钥
     * @param domain 向量化类型，可选值：query(用户问题向量化)、para(知识原文向量化)
     */
    public XingfengEmbeddingModel(String apiHost, String appId, String apiKey, String apiSecret, String domain) {
        this.appId = appId;
        this.apiKey = apiKey;
        this.apiSecret = apiSecret;
        this.apiHost = apiHost != null && !apiHost.isEmpty() ? apiHost : DEFAULT_API_HOST;
        this.domain = domain != null && !domain.isEmpty() ? domain : "query";
        
        // 配置OkHttpClient
        this.httpClient = new OkHttpClient.Builder()
                .connectTimeout(10, TimeUnit.SECONDS)
                .readTimeout(30, TimeUnit.SECONDS)
                .writeTimeout(30, TimeUnit.SECONDS)
                .build();
        
        log.info("初始化讯飞文本向量化API客户端: apiHost={}, domain={}", this.apiHost, this.domain);
    }

    @Override
    public Response<List<Embedding>> embedAll(List<TextSegment> textSegments) {
        List<Embedding> embeddings = new ArrayList<>();
        for (TextSegment segment : textSegments) {
            try {
                Embedding embedding = embedText(segment.text());
                embeddings.add(embedding);
            } catch (Exception e) {
                log.error("获取文本嵌入向量失败", e);
                throw new RuntimeException("获取文本嵌入向量失败: " + e.getMessage(), e);
            }
        }
        return Response.from(embeddings);
    }
    
    /**
     * 嵌入多个文本字符串
     */
    public Response<List<Embedding>> embedStrings(List<String> texts) {
        List<Embedding> embeddings = new ArrayList<>();
        for (String text : texts) {
            try {
                Embedding embedding = embedText(text);
                embeddings.add(embedding);
            } catch (Exception e) {
                log.error("获取文本嵌入向量失败", e);
                throw new RuntimeException("获取文本嵌入向量失败: " + e.getMessage(), e);
            }
        }
        return Response.from(embeddings);
    }

    private Embedding embedText(String text) {
        try {
            // 构建请求体
            JSONObject requestBody = buildRequestBody(text);
            
            // 获取鉴权URL
            String authUrl = getAuthUrl(apiHost, apiKey, apiSecret);
            
            log.debug("向量化API请求URL: {}", authUrl);
            log.debug("向量化API请求体: {}", requestBody.toJSONString());
            
            // 创建HTTP请求
            RequestBody body = RequestBody.create(
                    MediaType.parse("application/json"),
                    requestBody.toJSONString()
            );
            
            Request request = new Request.Builder()
                    .url(authUrl)
                    .post(body)
                    .build();

            // 发送请求并获取响应
            try (okhttp3.Response httpResponse = httpClient.newCall(request).execute()) {
                if (!httpResponse.isSuccessful()) {
                    String errorBody = httpResponse.body() != null ? httpResponse.body().string() : "无响应体";
                    log.error("文本向量化API调用失败: 状态码={}, 响应体={}", httpResponse.code(), errorBody);
                    throw new IOException("Unexpected code " + httpResponse);
                }
                
                String responseBody = httpResponse.body().string();
                log.debug("向量化API响应: {}", responseBody);
                
                JSONObject jsonResponse = JSON.parseObject(responseBody);

                // 解析响应
                if (jsonResponse.getJSONObject("header").getInteger("code") != 0) {
                    String message = jsonResponse.getJSONObject("header").getString("message");
                    throw new RuntimeException("文本向量化API错误: " + message);
                }
                
                // 获取向量
                JSONArray vector = jsonResponse.getJSONObject("payload").getJSONObject("vector").getJSONArray("float_vector");
                
                // 转换为float数组
                float[] floatVector = new float[vector.size()];
                for (int i = 0; i < vector.size(); i++) {
                    floatVector[i] = vector.getFloatValue(i);
                }
                
                return new Embedding(floatVector);
            }
        } catch (Exception e) {
            log.error("文本向量化API调用失败", e);
            throw new RuntimeException("文本向量化API调用失败: " + e.getMessage(), e);
        }
    }

    /**
     * 构建请求体
     */
    private JSONObject buildRequestBody(String text) {
        JSONObject requestBody = new JSONObject();
        
        // header部分
        JSONObject header = new JSONObject();
        header.put("app_id", appId);
        header.put("uid", UUID.randomUUID().toString().substring(0, 10));
        header.put("status", 3); // 一次传完
        requestBody.put("header", header);
        
        // parameter部分
        JSONObject parameter = new JSONObject();
        JSONObject emb = new JSONObject();
        emb.put("domain", domain); // query或para
        
        JSONObject feature = new JSONObject();
        feature.put("encoding", "utf8");
        feature.put("compress", "raw");
        feature.put("format", "plain");
        emb.put("feature", feature);
        
        parameter.put("emb", emb);
        requestBody.put("parameter", parameter);
        
        // payload部分
        JSONObject payload = new JSONObject();
        JSONObject messages = new JSONObject();
        messages.put("encoding", "utf8");
        messages.put("compress", "raw");
        messages.put("format", "json");
        messages.put("status", 3);
        messages.put("text", text);
        payload.put("messages", messages);
        requestBody.put("payload", payload);
        
        return requestBody;
    }
    
    /**
     * 获取认证之后的URL
     * 注意：讯飞文本向量化API的鉴权方式与大模型API不同
     */
    private String getAuthUrl(String apiUrl, String apiKey, String apiSecret) throws Exception {
        URL url = new URL(apiUrl);
        // 时间
        SimpleDateFormat format = new SimpleDateFormat("EEE, dd MMM yyyy HH:mm:ss z", Locale.US);
        format.setTimeZone(TimeZone.getTimeZone("GMT"));
        String date = format.format(new Date());
        
        // 拼接签名原文
        String preStr = "host: " + url.getHost() + "\n" +
                "date: " + date + "\n" +
                "POST " + url.getPath() + " HTTP/1.1";
        
        // SHA256加密
        Mac mac = Mac.getInstance("hmacsha256");
        SecretKeySpec spec = new SecretKeySpec(apiSecret.getBytes(StandardCharsets.UTF_8), "hmacsha256");
        mac.init(spec);

        byte[] hexDigits = mac.doFinal(preStr.getBytes(StandardCharsets.UTF_8));
        // Base64加密
        String sha = Base64.getEncoder().encodeToString(hexDigits);
        // 拼接
        String authorization = String.format("api_key=\"%s\", algorithm=\"%s\", headers=\"%s\", signature=\"%s\"", 
                apiKey, "hmac-sha256", "host date request-line", sha);
        
        // 构建完整的URL
        StringBuilder authUrlBuilder = new StringBuilder();
        authUrlBuilder.append(url.getProtocol()).append("://").append(url.getHost());
        
        // 添加路径
        if (url.getPath() != null && !url.getPath().isEmpty() && !url.getPath().equals("/")) {
            authUrlBuilder.append(url.getPath());
        }
        
        // 添加查询参数
        authUrlBuilder.append("?authorization=")
                .append(Base64.getEncoder().encodeToString(authorization.getBytes(StandardCharsets.UTF_8)))
                .append("&date=").append(java.net.URLEncoder.encode(date, "UTF-8"))
                .append("&host=").append(url.getHost());

        return authUrlBuilder.toString();
    }
} 